import numpy as np
import tensorflow as tf

# Superparameters
GAMMA = 0.9     # reward discount in TD error

#用于创建actor行动网络
class Actor(object):
    #初始化actor实例
    def __init__(self, sess, n_features, n_actions, lr=0.001):
        self.sess = sess

        self.s = tf.compat.v1.placeholder(tf.float32, [None, n_features], "state")#状态占位符，none表示可以接受不同数量的样本
        self.a = tf.compat.v1.placeholder(tf.int32, None, "act")#动作占位符，整数类型
        self.td_error = tf.compat.v1.placeholder(tf.float32, None, "td_error")  # TD_error（时序差分误差）：当前状态下采取动作后获得的奖励与预期奖励之间的误差



        #创建神经网络的变量作用域 "Actor"
        with tf.compat.v1.variable_scope('Actor'):
            #创建神经网络的第一层（隐藏层）使用ReLU激活函数
            l1 = tf.keras.layers.Dense(
                units=10,  # number of hidden units 输出单元数
                activation=tf.nn.relu,
                kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.3),  # weights
                bias_initializer=tf.constant_initializer(0.1),  # biases
                name='l1'
            )(self.s)
            #创建第二层神经网络
            self.acts_prob = tf.keras.layers.Dense(
                units=n_actions,  # output units 输出单元数，对应于可用动作数量
                activation=tf.nn.softmax,  # get action probabilities
                kernel_initializer=tf.keras.initializers.RandomNormal(0.0, 0.3),  # weights
                bias_initializer=tf.constant_initializer(0.1),  # biases
                name='acts_prob'
            )(l1)

        # 创建变量作用域 "exp_v"，计算 Actor 的损失函数（损失是根据策略梯度方法计算的优势函数）
        with tf.compat.v1.variable_scope('exp_v'):
            #计算选择的动作的概率的对数（log probability）并与 TD error 相乘，然后取平均值，这构成了优势函数（advantage）
            log_prob = tf.compat.v1.log(self.acts_prob[0, self.a])
            self.exp_v = tf.compat.v1.reduce_mean(log_prob * self.td_error)  # advantage (TD_error) guided loss

        #创建变量作用域 "train"，定义训练操作，使用 Adam 优化器来最大化优势函数
        with tf.compat.v1.variable_scope('train'):
            self.train_op = tf.compat.v1.train.AdamOptimizer(lr).minimize(-self.exp_v)  # minimize(-exp_v) = maximize(exp_v)

    #actor类中的函数:执行神经网络的训练步骤
    def learn(self, s, a, td):
        # s = s[np.newaxis, :]
        #创建一个dict字典，将输入数据提供给TensorFlow计算图
        feed_dict = {self.s: s, self.a: a, self.td_error: td}
        _, exp_v = self.sess.run([self.train_op, self.exp_v], feed_dict)
        return exp_v
    #run执行训练操作train和计算优势函数exp操作
    #run返回两个值，只关心exp_v（当前策略的优势函数值）用于评估选择的动作质量。exp_v被返回供调用者使用，监控和记录训练中策略的性能。

    #actor类中的函数：根据当前状态s选择一个动作。
    def choose_action(self, s):
        # s = s[np.newaxis, :]
        probs = self.sess.run(self.acts_prob, feed_dict={self.s: s})   # get probabilities for all actions获取所有可能的动作
        return np.random.choice(np.arange(probs.shape[1]), p=probs.ravel())   # return a int。根据概率分布选取一个动作

#创建Critic（评论者）神经网络：用于强化学习中的值函数估计
class Critic(object):
    def __init__(self, sess, n_features, lr=0.01):
        self.sess = sess

        self.s = tf.compat.v1.placeholder(tf.float32, [1, n_features], "state")#状态占位符，一次处理一个状态
        self.v_ = tf.compat.v1.placeholder(tf.float32, [1, 1], "v_next")#下一个状态的估计值的占位符
        self.r = tf.compat.v1.placeholder(tf.float32, None, 'r')#奖励（reward）占位符，浮点数

        #创建神经网络的变量作用域‘Critic’
        with tf.compat.v1.variable_scope('Critic'):
            #创建神经网络的第一层（隐藏层），使用ReLU激活函数
            l1 = tf.keras.layers.Dense(
                units=10,  # number of hidden units输出单元数
                activation=tf.nn.relu,  # None
                kernel_initializer=tf.keras.initializers.RandomNormal(0., .3),  # weights
                bias_initializer=tf.constant_initializer(0.1),  # biases
                name='l1'
            )(self.s)

            #创建神经网络的第二层，输出值函数的估计值
            self.v = tf.keras.layers.Dense(
                units=1,  # output units 输出单元数为1，因为Critic目标是估计状态值函数的值
                activation=None,
                kernel_initializer=tf.keras.initializers.RandomNormal(0., .3),  # weights
                bias_initializer=tf.constant_initializer(0.1),  # biases
                name='V'
            )(l1)

        #创建变量作用域‘squared_TD_error：计算TD error和损失函数
        with tf.compat.v1.variable_scope('squared_TD_error'):
            self.td_error = self.r + GAMMA * self.v_ - self.v  #GAMMA为常数，表示折扣因子，用于考虑未来奖励的影响。
            self.loss = tf.square(self.td_error)    # TD_error = (r+gamma*V_next) - V_eval
        
        #创建变量作用域’train‘：定义训练操作，使用adam优化器来最小化损失函数。
        with tf.compat.v1.variable_scope('train'):
            self.train_op = tf.compat.v1.train.AdamOptimizer(lr).minimize(self.loss)

    #Critic类中的函数：执行值函数的训练步骤
    def learn(self, s, r, s_):
        # s, s_ = s[np.newaxis, :], s_[np.newaxis, :]

        #run计算下一个状态’S_'对应的值函数估计值‘V_'。这里将状态占位符换成’s_‘来计算下一状态的值函数估计值
        v_ = self.sess.run(self.v, {self.s: s_})
        #计算TD error和执行Critic神经网络的训练操作
        td_error, _ = self.sess.run([self.td_error, self.train_op],
                                    {self.s: s, self.v_: v_, self.r: r})
        return td_error
        #td_error表示当前状态下的TD error值，是当前状态值函数估计的误差。在强化学习中，TD error 通常用于评估值函数的质量
        #td_error被返回，供调用者使用。值函数的训练目标是减小TD error，提高值函数的估计精度,更好的指导决策